"""
Takes in a directory of transition files and outputs statistics about transition probabilities
"""

import os 
import sys
import datetime


def main():
	inputdir = sys.argv[1]
	outputdir = sys.argv[2]

	clusters = int(sys.argv[3])

	inputFileList = os.listdir(inputdir)

	transitions = {}
	prob = {}
	post = {}
	for i in range(clusters + 1):
		transitions[i] = {}
		prob[i] = {}
		post[i] = {}
		for j in range(clusters + 1):
			transitions[i][j] = 0
			prob[i][j] = 0
			post[i][j] = 0

	for inputFileName in inputFileList:

		inputFilePath = inputdir+'/'+inputFileName

		if os.path.isfile(inputFilePath):
			processFile(inputFilePath, transitions)

	for i in range(clusters + 1):

		out = transitionsOutOf(transitions, i)
		into = transitionsInto(transitions, i)

		for j in range(clusters + 1):

			if out:
				prob[i][j] = float(transitions[i][j]) / float(out)
			if into:
				post[j][i] = float(transitions[j][i]) / float(into)


	countfile = open(outputdir + '/' + 'transition_count.stat', 'w')
	probfile = open(outputdir + '/' + 'transition_prob.stat', 'w')
	postfile = open(outputdir + '/' + 'transition_posterior.stat', 'w')

	for i in range(clusters + 1):

		for j in range(clusters + 1):
			countfile.write('\t'.join([str(i), str(j), str(transitions[i][j])]) + '\n')
			probfile.write('\t'.join([str(i), str(j), str(prob[i][j])]) + '\n')
			postfile.write('\t'.join([str(i), str(j), str(post[i][j])]) + '\n')

	countfile.close()
	probfile.close()
	postfile.close()


def transitionsOutOf(transitions, cluster):

	total = 0
	for other in transitions[cluster]:
		total += transitions[cluster][other]

	return total

def transitionsInto(transitions, cluster):

	total = 0
	for other in transitions:
		total += transitions[other][cluster]

	return total

def processFile(inputfilepath, transitions):

	inputfile = open(inputfilepath, 'r')

	for line in inputfile:

		fields = parseLine(line)

		transitions[fields['current']][fields['next']] += 1

		
	inputfile.close()

def parseLine(line):

	tokens = line.strip().split('\t')

	fields = {}
	fields['username'] = tokens[0]
	fields['cohort'] = int(tokens[1])
	fields['lifespan'] = int(tokens[2])
	fields['offset'] = int(tokens[3])
	fields['current'] = int(tokens[4])
	fields['next'] = int(tokens[5])

	return fields

main()
